package learning.crf.training;

import learning.crf.inference.IParseScorer;
import learning.crf.model.Model;
import learning.crf.training.CRF.DecodingResult;
import learning.data.BalancedResampler;
import learning.data.Dataset;
import learning.data.document.SequenceDocument;
import learning.logging.Log;
import learning.logging.Log.Level;

public class CollinsTraining {

	//public int maxIterations = 70;
	public int maxIterations = 50;
	public float avgIterationsProportion = 0.75f;
	public boolean computeAvgParameters = true;
		
	private IParseScorer allParseScorer;
	
	public CollinsTraining() {
		allParseScorer = new BasicParseScorer(BasicParseScorer.Type.ALL_TERMINALS);
	}
	
	public CollinsTraining(boolean learnTransitionFeatures) {
		allParseScorer = new BasicParseScorer(BasicParseScorer.Type.ALL_TERMINALS,
				learnTransitionFeatures);
	}
	
	public void crossValidation(Dataset<SequenceDocument> dataset, int c, EvaluatorCV evaluator) {
		evaluator.init(dataset.numLabels());
		
		for (int i=0; i < c; i++) {
			System.out.println("training cv set " + i);
			dataset.setIteratorView(c, i, true);
			//Dataset<SequenceDocument> balancedDataset = (new BalancedResampler()).resample(dataset);
			//CRFParameters parameters = train(balancedDataset);
			CRFParameters parameters = train(dataset);

			System.out.println("testing cv set " + i);
			dataset.setIteratorView(c, i, false);
			Evaluator e = new Evaluator(evaluator.level);
			test(dataset, parameters, e);
			evaluator.update(e, parameters);
			System.out.println(FormatHelper.formatConfusionMatrix(evaluator));
		}
	}	
	
	public void test(Dataset<SequenceDocument> testData, CRFParameters parameters, Evaluator evaluator) {
		parameters.model = new Model(testData.numLabels());

		allParseScorer.setDataset(testData);
		
		evaluator.init(testData.numLabels());

		for (SequenceDocument doc : testData) {
			// compute most likely label under current parameters
			DecodingResult predictedParse = CRF.decode(doc, allParseScorer, parameters);
			DecodingResult trueParse = new DecodingResult(doc.labels, 0);
			
			// update evaluator
			evaluator.update(doc, predictedParse, trueParse);
		}
	}
	
	public CRFParameters train(Dataset<SequenceDocument> trainingData) {
		
		allParseScorer.setDataset(trainingData);
		
		CRFParameters avgParameters = new CRFParameters();
		CRFParameters iterAvgParameters = new CRFParameters();
		CRFParameters iterParameters = new CRFParameters();
		
		avgParameters.model = iterAvgParameters.model = iterParameters.model = new Model(trainingData.numLabels());
		avgParameters.reset();
		iterParameters.reset();
		
		int avgIterations = maxIterations - (int)(maxIterations*avgIterationsProportion);
		
		Evaluator evaluator = null;
		if (Level.FINE.equals(Log.getLevel()))
			evaluator = new Evaluator(Evaluator.Level.COARSE, avgParameters.model.numStates);
		
		for (int i = 0; i < maxIterations; i++)
		{
			//System.out.println("iteration " + i);
			if (Level.FINE.equals(Log.getLevel()))
				evaluator.reset();
			
			//float delta = 1;
			float delta = (float)(2.0 / (2.0 + i));
			//float delta = (float)(20.0 / (20.0 + i));
			boolean useIterAverage = computeAvgParameters && (i >= maxIterations - avgIterations);
			
			trainingIteration(trainingData, delta, evaluator, iterParameters, useIterAverage, iterAvgParameters);

			if (computeAvgParameters && useIterAverage) {
				avgParameters.sum(iterAvgParameters, 1.0f / avgIterations);
			}
			
			if (Level.FINE.equals(Log.getLevel()))
				Log.log(FormatHelper.formatConfusionMatrix(evaluator), Level.FINE);
		}
		
		if (Level.FINE.equals(Log.getLevel()))
			Log.log(FormatHelper.formatConfusionMatrix(evaluator), Level.FINE);

		if (Level.FINE.equals(Log.getLevel()))
			Log.log(FormatHelper.formatParams(iterParameters, trainingData), Level.FINE);

		return (computeAvgParameters)? avgParameters : iterParameters;
	}
	
	public void trainingIteration(Dataset<SequenceDocument> trainingData, float delta, Evaluator evaluator,
			CRFParameters iterParameters, boolean useIterAverage, CRFParameters iterAvgParameters) {
		
		// resetting parameters to 0
		iterAvgParameters.reset();
		
		trainingData.randomizeOrder();
		
		for (SequenceDocument doc : trainingData) {
			
			// compute most likely label under current parameters			
			DecodingResult predictedParse = CRF.decode(doc, allParseScorer, iterParameters);
			DecodingResult trueParse = new DecodingResult(doc.labels, 0);
			
			// update weights
			CRF.update(doc, allParseScorer, predictedParse, trueParse, iterParameters, delta);
			
			//DebugInfo.print(iterParameters.transitionParameters[0][1], trainingData);

			//DebugInfo.print(doc, predictedParse, trueParse);
			
			if (useIterAverage) {				
				iterAvgParameters.sum(iterParameters, 1.0f / trainingData.numDocs());
			}			
			
			// update evaluator
			if (Level.FINE.equals(Log.getLevel()))
				evaluator.update(doc, predictedParse, trueParse);
		}
	}
}
